# -*- coding: UTF-8 -*-
import torch
from res2net import res2net50_48w_2s

model = res2net50_48w_2s(pretrained=False)
model.eval()
checkpoint = torch.load('res2net50_48w_2s-afed724a.pth', map_location='cpu')
model.load_state_dict(checkpoint)
batch_size = 1
input = torch.randn(batch_size, 3, 224, 224, requires_grad=True)
torch.onnx.export(model, input, 'res2net50_48w_2s.onnx', export_params=True,
                  opset_version=10, do_constant_folding=True, input_names=['input'], output_names=['output'],
                  dynamic_axes={'input': {0: '?'}, 'output': {0: '?'}})
